# @Time    : 2023/12/21 16:38
# @Author  : yinhai
# @File    : tasks.py
# @Project : ai_studio

import zipfile
from pathlib import Path

import cv2
import numpy as np
from ultralytics import YOLO, SAM
import json

from studio_celery.main import app


def zip(file_path, zip_filename=None):
    """
    压缩文件或文件夹
    :param zip_filename:
    :param file_path:
    :return:
    """
    if not zip_filename:
        zip_filename = str(file_path) + '.zip'
    with zipfile.ZipFile(zip_filename, 'w') as zip_obj:
        path = Path(file_path)
        if path.is_file():
            print(path, 'add zip')
            zip_obj.write(path, arcname=path.relative_to(path.parent))
        else:
            for i in path.rglob('*'):
                print(i, 'add zip')
                zip_obj.write(i, arcname=i.relative_to(path))
    return zip_filename


def unzip(zip_filename, unzip_path=None):
    zip_filename = Path(zip_filename)
    if not unzip_path:
        unzip_path = zip_filename.parent / zip_filename.stem
    with zipfile.ZipFile(zip_filename, "r") as zip_obj:
        zip_obj.extractall(path=unzip_path)

    return unzip_path


class Annotation:
    def __init__(self):
        self.SAM = SAM('/data/ai_studio/studio_celery/aicv/sam_l.pt')
        self.YOLO = YOLO('/data/ai_studio/studio_celery/aicv/yolov8x.pt')
        self.color = {
            "person": (255, 0, 0),  # Red
            "car": (0, 255, 0),  # Green
            "truck": (0, 0, 255),  # Blue
            "stop sign": (255, 255, 0),  # Yellow
            "traffic sign": (255, 255, 0),  # Yellow
            "motorcycle": (255, 0, 255),  # Magenta
            "traffic light": (0, 255, 255),  # Cyan
            "bicycle": (128, 0, 128),  # Purple
            "bus": (255, 165, 0),  # Orange,

            "人": (255, 0, 0),  # Red
            "汽车": (0, 255, 0),  # Green
            "卡车": (0, 0, 255),  # Blue
            "停车标志": (255, 255, 0),  # Yellow
            "交通标志": (255, 255, 0),  # Yellow
            "摩托车": (255, 0, 255),  # Magenta
            "红绿灯": (0, 255, 255),  # Cyan
            "单车": (128, 0, 128),  # Purple
            "公共汽车": (255, 165, 0)  # Orange

        }
        # for i in {'person', 'bicycle', 'car', 'motorcycle', 'bus', 'train', 'truck',
        #           'traffic light', 'stop sign'}:
        #     self.color[i] = tuple(np.random.randint(0, 255, 3).tolist())

    def run(self, path: Path, dirs: Path):
        img = cv2.imread(str(path))
        mask_img = np.zeros_like(img)
        eff_img = np.zeros_like(img)
        result = self.YOLO(source=str(path))
        model_label = result[0].names
        cls_label = result[0].boxes.cls.cpu().numpy()
        bbox = result[0].boxes.xyxy.cpu().numpy()
        json_data = []
        for index, item in enumerate(bbox):
            if (label := model_label[cls_label[index]]) not in {'person', 'bicycle', 'car', 'motorcycle', 'bus',
                                                                'truck',
                                                                'traffic light', 'stop sign'}: continue
            sam_result = self.SAM(img, bboxes=item)
            mask = sam_result[0].masks.xy
            json_data.append({'label': label, 'bbox': item.tolist()})
            for mask_idx in mask:
                cv2.fillPoly(mask_img, [mask_idx.astype(np.int32)], (255, 255, 255))
                cv2.fillPoly(eff_img, [mask_idx.astype(np.int32)], self.color[label])
                json_data.append({'label': label, 'mask': mask_idx.tolist()})

        save_path_mask = dirs.parent / (dirs.stem + '_export') / path.relative_to(dirs).parent / (path.stem + '.png')
        save_path_eff = dirs.parent / (dirs.stem + '_export') / path.relative_to(dirs).parent / (path.stem + '.jpg')
        save_path_json = dirs.parent / (dirs.stem + '_export') / path.relative_to(dirs).parent / (path.stem + '.json')
        save_path_eff.parent.mkdir(exist_ok=True, parents=True)
        cv2.imwrite(str(save_path_mask), mask_img)
        cv2.imwrite(str(save_path_eff), cv2.addWeighted(img, 0.8, eff_img, 0.2, 0.0))
        with save_path_json.open('w', encoding='utf-8') as f:
            f.write(json.dumps(json_data, ensure_ascii=False, indent=4))

    def runData(self, json_path: Path, dirs: Path):
        with json_path.open('r', encoding='utf-8') as f:
            json_data = json.loads(f.read())
        path = json_path.with_suffix('.jpg')
        if not path.exists():
            return

        img = cv2.imread(str(path))
        mask_img = np.zeros_like(img)
        eff_img = np.zeros_like(img)
        all_json_data = []
        for index, item in enumerate(json_data):
            if 'bbox' not in item:
                continue
            print(item, json_path)
            label = item['label']
            all_json_data.append(item)
            sam_result = self.SAM(img, bboxes=item['bbox'])
            mask = sam_result[0].masks.xy
            for mask_idx in mask:
                cv2.fillPoly(mask_img, [mask_idx.astype(np.int32)], (255, 255, 255))
                cv2.fillPoly(eff_img, [mask_idx.astype(np.int32)], self.color[label])
                all_json_data.append({'label': label, 'mask': mask_idx.tolist(), 'property': item['property']})

        save_path_mask = dirs.parent / (dirs.stem + '_export') / path.relative_to(dirs).parent / (path.stem + '.png')
        save_path_eff = dirs.parent / (dirs.stem + '_export') / path.relative_to(dirs).parent / (path.stem + '.jpg')
        save_path_json = dirs.parent / (dirs.stem + '_export') / path.relative_to(dirs).parent / (path.stem + '.json')
        save_path_eff.parent.mkdir(exist_ok=True, parents=True)
        cv2.imwrite(str(save_path_mask), mask_img)
        cv2.imwrite(str(save_path_eff), cv2.addWeighted(img, 0.8, eff_img, 0.2, 0.0))
        with save_path_json.open('w', encoding='utf-8') as f:
            f.write(json.dumps(all_json_data, ensure_ascii=False, indent=4))


@app.task(bind=True)
def aicv(self, filename):
    anno = Annotation()
    unzip_path = unzip(filename)
    files = list(Path(unzip_path).rglob('*.json'))
    if files:
        for index, i in enumerate(files):
            anno.runData(i, unzip_path)
            self.update_state(state='PROGRESS',
                              meta={'current': index + 1, 'total': len(files)})
    else:
        files = list(Path(unzip_path).rglob('*.jpg')) + list(Path(unzip_path).rglob('*.png'))
        for index, img in enumerate(files):
            anno.run(img, unzip_path)
            self.update_state(state='PROGRESS',
                              meta={'current': index + 1, 'total': len(files)})

    if not (unzip_path.parent / (unzip_path.stem + '_export')).exists():
        self.update_state(state='ERROR')
        unzip_path.unlink(missing_ok=True)
        return 'ERROR'

    zip_file = zip(unzip_path.parent / (unzip_path.stem + '_export'))

    return str(zip_file)
